"""
基础策略类
提供策略开发的基础框架和通用功能
"""

import pandas as pd
import numpy as np
from abc import ABC, abstractmethod
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Callable
import logging
from dataclasses import dataclass
from enum import Enum


class SignalType(Enum):
    """交易信号类型"""
    LONG = "long"      # 做多
    SHORT = "short"    # 做空
    CLOSE = "close"    # 平仓
    HOLD = "hold"      # 持有


class StrategyState(Enum):
    """策略状态"""
    INITIALIZED = "initialized"    # 已初始化
    RUNNING = "running"           # 运行中
    PAUSED = "paused"            # 暂停
    STOPPED = "stopped"          # 已停止
    ERROR = "error"              # 错误状态


@dataclass
class TradingSignal:
    """交易信号数据类"""
    strategy_name: str
    symbol: str
    signal_type: SignalType
    volume: int
    price: float
    timestamp: datetime
    confidence: float = 1.0
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


@dataclass
class PerformanceMetrics:
    """策略性能指标"""
    total_return: float = 0.0
    annual_return: float = 0.0
    sharpe_ratio: float = 0.0
    max_drawdown: float = 0.0
    win_rate: float = 0.0
    profit_factor: float = 0.0
    total_trades: int = 0
    winning_trades: int = 0
    losing_trades: int = 0
    avg_win: float = 0.0
    avg_loss: float = 0.0
    largest_win: float = 0.0
    largest_loss: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'total_return': self.total_return,
            'annual_return': self.annual_return,
            'sharpe_ratio': self.sharpe_ratio,
            'max_drawdown': self.max_drawdown,
            'win_rate': self.win_rate,
            'profit_factor': self.profit_factor,
            'total_trades': self.total_trades,
            'winning_trades': self.winning_trades,
            'losing_trades': self.losing_trades,
            'avg_win': self.avg_win,
            'avg_loss': self.avg_loss,
            'largest_win': self.largest_win,
            'largest_loss': self.largest_loss
        }


class BaseStrategy(ABC):
    """
    基础策略类
    
    所有策略都应该继承此类并实现抽象方法
    提供策略生命周期管理、数据处理、信号生成等基础功能
    """
    
    def __init__(self, name: str, config: Dict[str, Any], data_manager: Any):
        """
        初始化基础策略
        
        Args:
            name: 策略名称
            config: 策略配置参数
            data_manager: 数据管理器实例
        """
        self.name = name
        self.config = config
        self.data_manager = data_manager
        
        # 策略状态
        self.state = StrategyState.INITIALIZED
        self.start_time = None
        self.stop_time = None
        
        # 交易相关
        self.symbol = config.get('symbol', '')
        self.position = 0  # 当前持仓
        self.entry_price = 0.0  # 入场价格
        self.last_signal = None
        
        # 性能统计
        self.trades = []  # 交易记录
        self.equity_curve = []  # 权益曲线
        self.performance_metrics = PerformanceMetrics()
        
        # 风险控制参数
        self.max_position = config.get('max_position', 1)
        self.stop_loss_pct = config.get('stop_loss_pct', 0.02)  # 2%止损
        self.take_profit_pct = config.get('take_profit_pct', 0.04)  # 4%止盈
        
        # 数据缓存
        self.kline_data = pd.DataFrame()
        self.current_bar = None
        self.current_tick = None
        
        # 回调函数
        self.signal_callback = None
        self.error_callback = None
        
        # 日志记录
        self.logger = logging.getLogger(f"strategy.{name}")
        
        self.logger.info(f"策略初始化完成: {name}")
    
    def initialize(self) -> None:
        """
        策略初始化
        
        在策略开始运行前调用，用于加载历史数据、初始化指标等
        子类可以重写此方法实现自定义初始化逻辑
        """
        try:
            self.logger.info(f"开始初始化策略: {self.name}")
            
            # 加载历史数据
            self._load_historical_data()
            
            # 调用子类的初始化方法
            self.on_initialize()
            
            # 注册数据回调
            if self.data_manager:
                self.data_manager.register_strategy_callback(
                    self.name, 
                    [self.symbol], 
                    self._on_data_update
                )
            
            self.state = StrategyState.INITIALIZED
            self.logger.info(f"策略初始化成功: {self.name}")
            
        except Exception as e:
            self.state = StrategyState.ERROR
            self.logger.error(f"策略初始化失败: {self.name}, 错误: {e}")
            if self.error_callback:
                self.error_callback(self.name, f"初始化失败: {e}")
            raise
    
    def start(self) -> None:
        """启动策略"""
        try:
            if self.state != StrategyState.INITIALIZED:
                raise ValueError(f"策略状态错误，无法启动: {self.state}")
            
            self.start_time = datetime.now()
            self.state = StrategyState.RUNNING
            
            # 调用子类的启动方法
            self.on_start()
            
            self.logger.info(f"策略已启动: {self.name}")
            
        except Exception as e:
            self.state = StrategyState.ERROR
            self.logger.error(f"策略启动失败: {self.name}, 错误: {e}")
            if self.error_callback:
                self.error_callback(self.name, f"启动失败: {e}")
            raise
    
    def stop(self) -> None:
        """停止策略"""
        try:
            if self.state not in [StrategyState.RUNNING, StrategyState.PAUSED]:
                self.logger.warning(f"策略状态错误，无法停止: {self.state}")
                return
            
            self.stop_time = datetime.now()
            self.state = StrategyState.STOPPED
            
            # 注销数据回调
            if self.data_manager:
                self.data_manager.unregister_strategy_callback(self.name)
            
            # 调用子类的停止方法
            self.on_stop()
            
            # 计算最终性能指标
            self._calculate_performance_metrics()
            
            self.logger.info(f"策略已停止: {self.name}")
            
        except Exception as e:
            self.state = StrategyState.ERROR
            self.logger.error(f"策略停止失败: {self.name}, 错误: {e}")
            if self.error_callback:
                self.error_callback(self.name, f"停止失败: {e}")
    
    def pause(self) -> None:
        """暂停策略"""
        if self.state == StrategyState.RUNNING:
            self.state = StrategyState.PAUSED
            self.on_pause()
            self.logger.info(f"策略已暂停: {self.name}")
    
    def resume(self) -> None:
        """恢复策略"""
        if self.state == StrategyState.PAUSED:
            self.state = StrategyState.RUNNING
            self.on_resume()
            self.logger.info(f"策略已恢复: {self.name}")
    
    def cleanup(self) -> None:
        """
        策略清理
        
        释放资源，保存状态等
        """
        try:
            # 注销数据回调
            if self.data_manager:
                self.data_manager.unregister_strategy_callback(self.name)
            
            # 调用子类的清理方法
            self.on_cleanup()
            
            self.logger.info(f"策略清理完成: {self.name}")
            
        except Exception as e:
            self.logger.error(f"策略清理失败: {self.name}, 错误: {e}")
    
    def _on_data_update(self, symbol: str, data: Dict[str, Any]) -> None:
        """
        数据更新回调
        
        Args:
            symbol: 合约代码
            data: 数据字典
        """
        try:
            if self.state != StrategyState.RUNNING:
                return
            
            if symbol != self.symbol:
                return
            
            # 更新当前数据
            self.current_tick = data
            
            # 调用tick数据处理
            self.on_tick(data)
            
            # 检查是否有新的K线数据
            if self._is_new_bar(data):
                self._update_kline_data(data)
                self.on_bar(self.current_bar)
            
        except Exception as e:
            self.logger.error(f"数据处理失败: {self.name}, 错误: {e}")
            if self.error_callback:
                self.error_callback(self.name, f"数据处理失败: {e}")
    
    def _load_historical_data(self) -> None:
        """加载历史数据"""
        try:
            if not self.data_manager:
                return
            
            # 获取配置的数据长度
            data_length = self.config.get('data_length', 200)
            duration = self.config.get('duration', 60)  # 默认1分钟K线
            
            # 加载K线数据
            self.kline_data = self.data_manager.get_kline_data(
                self.symbol, duration, data_length
            )
            
            if not self.kline_data.empty:
                self.current_bar = self.kline_data.iloc[-1].to_dict()
                self.logger.info(f"加载历史数据成功: {self.symbol}, 数据量: {len(self.kline_data)}")
            else:
                self.logger.warning(f"未获取到历史数据: {self.symbol}")
                
        except Exception as e:
            self.logger.error(f"加载历史数据失败: {self.symbol}, 错误: {e}")
    
    def _is_new_bar(self, tick_data: Dict[str, Any]) -> bool:
        """
        判断是否有新的K线数据
        
        Args:
            tick_data: tick数据
            
        Returns:
            是否有新K线
        """
        if self.current_bar is None:
            return True
        
        # 简单的时间判断，实际实现可能需要更复杂的逻辑
        current_time = tick_data.get('datetime')
        if current_time is None:
            return False
        
        # 这里需要根据K线周期来判断
        # 暂时简化处理
        return False
    
    def _update_kline_data(self, tick_data: Dict[str, Any]) -> None:
        """
        更新K线数据
        
        Args:
            tick_data: tick数据
        """
        # 这里应该根据tick数据更新K线
        # 暂时简化处理
        pass
    
    def generate_signal(self) -> Optional[TradingSignal]:
        """
        生成交易信号
        
        Returns:
            交易信号，如果没有信号则返回None
        """
        try:
            if self.state != StrategyState.RUNNING:
                return None
            
            # 调用子类的信号生成方法
            signal = self.on_signal()
            
            if signal:
                # 验证信号
                if self._validate_signal(signal):
                    self.last_signal = signal
                    
                    # 调用信号回调
                    if self.signal_callback:
                        self.signal_callback(signal)
                    
                    self.logger.info(f"生成交易信号: {signal.signal_type.value}, "
                                   f"价格: {signal.price}, 数量: {signal.volume}")
                    return signal
                else:
                    self.logger.warning(f"信号验证失败: {signal}")
            
            return None
            
        except Exception as e:
            self.logger.error(f"信号生成失败: {self.name}, 错误: {e}")
            return None
    
    def calculate_position_size(self, signal: TradingSignal) -> int:
        """
        计算仓位大小
        
        Args:
            signal: 交易信号
            
        Returns:
            仓位大小
        """
        try:
            # 调用子类的仓位计算方法
            position_size = self.on_position_sizing(signal)
            
            # 应用风险控制
            position_size = self._apply_risk_control(position_size, signal)
            
            return position_size
            
        except Exception as e:
            self.logger.error(f"仓位计算失败: {self.name}, 错误: {e}")
            return 0
    
    def _validate_signal(self, signal: TradingSignal) -> bool:
        """
        验证交易信号
        
        Args:
            signal: 交易信号
            
        Returns:
            信号是否有效
        """
        try:
            # 基本验证
            if not signal.symbol or signal.volume <= 0:
                return False
            
            if signal.price <= 0:
                return False
            
            # 仓位限制检查
            if signal.signal_type in [SignalType.LONG, SignalType.SHORT]:
                if abs(self.position + signal.volume) > self.max_position:
                    self.logger.warning(f"超出最大仓位限制: {self.max_position}")
                    return False
            
            return True
            
        except Exception as e:
            self.logger.error(f"信号验证失败: {e}")
            return False
    
    def _apply_risk_control(self, position_size: int, signal: TradingSignal) -> int:
        """
        应用风险控制
        
        Args:
            position_size: 原始仓位大小
            signal: 交易信号
            
        Returns:
            风控后的仓位大小
        """
        # 最大仓位限制
        if signal.signal_type in [SignalType.LONG, SignalType.SHORT]:
            max_allowed = self.max_position - abs(self.position)
            position_size = min(position_size, max_allowed)
        
        return max(0, position_size)
    
    def _calculate_performance_metrics(self) -> None:
        """计算性能指标"""
        try:
            if not self.trades:
                return
            
            # 计算基本统计
            returns = [trade.get('pnl', 0) for trade in self.trades]
            winning_trades = [r for r in returns if r > 0]
            losing_trades = [r for r in returns if r < 0]
            
            self.performance_metrics.total_trades = len(self.trades)
            self.performance_metrics.winning_trades = len(winning_trades)
            self.performance_metrics.losing_trades = len(losing_trades)
            
            if self.performance_metrics.total_trades > 0:
                self.performance_metrics.win_rate = (
                    self.performance_metrics.winning_trades / self.performance_metrics.total_trades
                )
            
            # 计算收益指标
            total_pnl = sum(returns)
            self.performance_metrics.total_return = total_pnl
            
            if winning_trades:
                self.performance_metrics.avg_win = np.mean(winning_trades)
                self.performance_metrics.largest_win = max(winning_trades)
            
            if losing_trades:
                self.performance_metrics.avg_loss = np.mean(losing_trades)
                self.performance_metrics.largest_loss = min(losing_trades)
            
            # 计算盈亏比
            if self.performance_metrics.avg_loss != 0:
                self.performance_metrics.profit_factor = (
                    abs(self.performance_metrics.avg_win) / abs(self.performance_metrics.avg_loss)
                )
            
            self.logger.info(f"性能指标计算完成: {self.name}")
            
        except Exception as e:
            self.logger.error(f"性能指标计算失败: {self.name}, 错误: {e}")
    
    def get_performance_metrics(self) -> Dict[str, Any]:
        """
        获取策略性能指标
        
        Returns:
            性能指标字典
        """
        # 实时更新性能指标
        self._calculate_performance_metrics()
        
        metrics = self.performance_metrics.to_dict()
        metrics.update({
            'strategy_name': self.name,
            'symbol': self.symbol,
            'state': self.state.value,
            'start_time': self.start_time,
            'stop_time': self.stop_time,
            'current_position': self.position,
            'last_signal': self.last_signal.signal_type.value if self.last_signal else None
        })
        
        return metrics
    
    def get_status(self) -> Dict[str, Any]:
        """
        获取策略状态信息
        
        Returns:
            状态信息字典
        """
        return {
            'name': self.name,
            'state': self.state.value,
            'symbol': self.symbol,
            'position': self.position,
            'start_time': self.start_time,
            'stop_time': self.stop_time,
            'last_signal': self.last_signal.signal_type.value if self.last_signal else None,
            'total_trades': len(self.trades),
            'config': self.config
        }
    
    def set_signal_callback(self, callback: Callable[[TradingSignal], None]) -> None:
        """设置信号回调函数"""
        self.signal_callback = callback
    
    def set_error_callback(self, callback: Callable[[str, str], None]) -> None:
        """设置错误回调函数"""
        self.error_callback = callback
    
    # 抽象方法 - 子类必须实现
    
    @abstractmethod
    def on_initialize(self) -> None:
        """策略初始化回调 - 子类实现"""
        pass
    
    @abstractmethod
    def on_bar(self, bar_data: Dict[str, Any]) -> None:
        """
        K线数据回调 - 子类实现
        
        Args:
            bar_data: K线数据字典
        """
        pass
    
    @abstractmethod
    def on_tick(self, tick_data: Dict[str, Any]) -> None:
        """
        Tick数据回调 - 子类实现
        
        Args:
            tick_data: Tick数据字典
        """
        pass
    
    @abstractmethod
    def on_signal(self) -> Optional[TradingSignal]:
        """
        信号生成回调 - 子类实现
        
        Returns:
            交易信号或None
        """
        pass
    
    @abstractmethod
    def on_position_sizing(self, signal: TradingSignal) -> int:
        """
        仓位计算回调 - 子类实现
        
        Args:
            signal: 交易信号
            
        Returns:
            仓位大小
        """
        pass
    
    # 可选的生命周期回调方法
    
    def on_start(self) -> None:
        """策略启动回调 - 子类可选实现"""
        pass
    
    def on_stop(self) -> None:
        """策略停止回调 - 子类可选实现"""
        pass
    
    def on_pause(self) -> None:
        """策略暂停回调 - 子类可选实现"""
        pass
    
    def on_resume(self) -> None:
        """策略恢复回调 - 子类可选实现"""
        pass
    
    def on_cleanup(self) -> None:
        """策略清理回调 - 子类可选实现"""
        pass